#include "StdAfx.h"
#include "MyNetworkAnalyzer.h"


MyNetworkAnalyzer::MyNetworkAnalyzer(void)
:myFrameParser(INVALID_HANDLE_VALUE)
,myCaptureFile(INVALID_HANDLE_VALUE)
{
	
}

MyNetworkAnalyzer::~MyNetworkAnalyzer(void)
{
}


void __stdcall
MyParserBuild(PVOID Context, ULONG StatusCode, LPCWSTR lpDescription, ULONG ErrorType)
{
	TRACE(L"%s\n", lpDescription);
}


HANDLE MyNetworkAnalyzer::MyLoadNPL(TCHAR* tszDisplayFilter,TCHAR* tszConversationFilter)
{
	myFrameParser= INVALID_HANDLE_VALUE;

	CString strError;

	ULONG ret;
	// Use NULL to load default NPL set. 
	do
	{
		ret = NmLoadNplParser(NULL, NmAppendRegisteredNplSets, MyParserBuild, 0, &myNplParser);
		if( ret != ERROR_SUCCESS)
		{
			TRACE("\nNmLoadNplParser Error:%02x",ret );
			break;
		}

		ret = NmCreateFrameParserConfiguration(myNplParser, MyParserBuild, 0, &myFrameParserConfig);
		if( ret != ERROR_SUCCESS)
		{
			TRACE("\nNmCreateFrameParserConfiguration Error:%02x",ret );
			NmCloseHandle(myNplParser);
			return INVALID_HANDLE_VALUE;
		}
	
				//Add filter
		ret = NmAddFilter(myFrameParserConfig, tszDisplayFilter, &myDisplayFilterID);
		if(ret != ERROR_SUCCESS)
		{
			TRACE( _T("Fail to load Add display fitler, error: \n"), ret);
		}

		ret = NmAddFilter(myFrameParserConfig, tszConversationFilter, &myConversationFilterID);
		if(ret != ERROR_SUCCESS)
		{
			TRACE( _T("Fail to load Add conversation fitler, error: \n"), ret);
		}
		

		ret = NmAddProperty(myFrameParserConfig, L"Property.TCPPayloadLength", &myTCPPayloadLengthID);
		if (ret != 0)
		{
			TRACE( _T("Failed to add Property.TCPPayloadLength, error 0x%X\n"), ret);			
		}
	
		ret = NmAddField( myFrameParserConfig, L"TCP.DstPort", &myTCPDstPortID);
		ret = NmAddField( myFrameParserConfig, L"TCP.SrcPort", &myTCPSrcPortID);
		if (ret != ERROR_SUCCESS)
		{				
			TRACE( _T("Failed to add field, TCP.SrcPort, error 0x%X\n"), ret);
		}

		ret = NmAddField( myFrameParserConfig, L"ipv4.SourceAddress", &myIPv4SrcAddress);
		ret = NmAddField( myFrameParserConfig, L"ipv4.DestinationAddress", &myIPv4DstAddress);


		ret = NmAddField( myFrameParserConfig, L"TCP.DataOffset.DataOffset", &myTCPDataOffsetID);
		if (ret != ERROR_SUCCESS)
		{
			TRACE( _T("Failed to add field, TCP.DataOffset, error 0x%X\n"), ret);
		}

		ret = NmCreateFrameParser( myFrameParserConfig, &myFrameParser);
		if (ret != ERROR_SUCCESS)
		{

			TRACE( _T("Failed to create frame parser, error 0x%X\n"), ret);
			NmCloseHandle(myFrameParserConfig);
			NmCloseHandle(myNplParser);
			return INVALID_HANDLE_VALUE;
		}
	}
	while(0);
	return myFrameParser;
}


void MyNetworkAnalyzer::GetFramePayload(HANDLE myParsedFrame,HANDLE myRawFrame,TDumps& dumps)
{
	ULONG ret;    
	UINT32 PayloadLen = 0;
	ULONG retlen;    
	NmPropertyValueType PropType;
	UINT8 TCPHeaderSize;
	ULONG TCPSrcOffset, TCPSrcSize;
	ULONG TCPSrcPort,TCPDstPort;
	ULONG IPv4SrcAddress,IPv4DstAddress;
	CString strError;

	// Get Payload Length    
	ret = NmGetPropertyValueById(	myFrameParser, myTCPPayloadLengthID, sizeof(PayloadLen), (PBYTE)&PayloadLen, &retlen, &PropType);
	if (ret != ERROR_SUCCESS)
	{
		strError.Format( _T("Error retrieving TCP Payload Length Property, err=%d\n"), ret);
		MessageBox( NULL,strError,_T("Error!"),0);
		return;    
	}

	if (PayloadLen > 0)
	{
		// Allocate a buffer based on the Payload Length Property.
		CDumpInfo* info = (CDumpInfo* )new char[ sizeof(CDumpInfo) + PayloadLen];		

		// Get the Data Offset, used to determine the TCP header size        
		ret = NmGetFieldValueNumber8Bit( myParsedFrame, myTCPDataOffsetID, &TCPHeaderSize);
		if (ret != ERROR_SUCCESS)
		{
			strError.Format( _T("Error retrieving TCP Header Length Field, err=%d\n"), ret);
			MessageBox( NULL,strError,_T("Error!"),0);
			return;        
		}

		// Get the Offset of TCP.SrcPort which is the first field in TCP.        
		ret = NmGetFieldOffsetAndSize( myParsedFrame, myTCPSrcPortID, &TCPSrcOffset, &TCPSrcSize);
		if (ret != ERROR_SUCCESS)
		{
			strError.Format( _T("Error retrieving TCP SRC Header/Offset, err=%d\n"), ret);
			MessageBox( NULL,strError,_T("Error!"),0);
			return;        
		}

		ret = NmGetFieldValueNumber32Bit(myParsedFrame, myTCPSrcPortID,(PUINT32) &TCPSrcPort);
		ret = NmGetFieldValueNumber32Bit(myParsedFrame, myTCPDstPortID, (PUINT32) &TCPDstPort);
		ret = NmGetFieldValueNumber32Bit(myParsedFrame, myIPv4SrcAddress, (PUINT32) &IPv4SrcAddress);
		ret = NmGetFieldValueNumber32Bit(myParsedFrame, myIPv4DstAddress, (PUINT32) &IPv4DstAddress);

		info->dstIP= IPv4DstAddress;
		info->srcIP= IPv4SrcAddress;
		info->dstPort = TCPDstPort;
		info->srcPort = TCPSrcPort;
		//strError.Format( _T("Offset: %d, Length: %d, HeaderLen: %d\n"), TCPSrcOffset/8, PayloadLen, TCPHeaderSize*4);
		//MessageBox( NULL,strError,_T("Info!"),0);


		// Read in the partial frame.  The Offset is in bits.  TCPHeaderSize is off by a factor of 4.     
		ret = NmGetPartialRawFrame(myRawFrame, TCPSrcOffset/8 + TCPHeaderSize*4, PayloadLen, (PBYTE)info->data, &retlen);
		info->dataLengh=retlen;		
		// Do what ever you want with buf now.  I'll assume it's ASCII and print it.
		dumps.push_back(info);
	}

}


BOOL MyNetworkAnalyzer::Init(const TCHAR *tszCapFile,TCHAR* tszDisplayFilter,TCHAR* tszConversationFilter,TDumps& dumps)
{

    ULONG ret = ERROR_SUCCESS;

		// Initialize the parser engine and return a frame parser.
	myFrameParser = MyLoadNPL(tszDisplayFilter,tszConversationFilter);
	if (myFrameParser == INVALID_HANDLE_VALUE)
	{
		AfxMessageBox( _T("Errors creating frame parser\n") );
		return FALSE;
	}

	if(ERROR_SUCCESS == NmOpenCaptureFile( tszCapFile , &myCaptureFile))
	{
		ULONG myFrameCount = 0;
		ret = NmGetFrameCount(myCaptureFile, &myFrameCount); 
		if(ret == ERROR_SUCCESS)
		{
			HANDLE myRawFrame = INVALID_HANDLE_VALUE;
			for(ULONG i = 0; i < myFrameCount; i++)
			{
				HANDLE myParsedFrame = INVALID_HANDLE_VALUE;

				ret = NmGetFrame(myCaptureFile, i, &myRawFrame);
				if(ret == ERROR_SUCCESS)
				{
					// The last parameter is for API to return reassembled frame if enabled 
					// NULL means that API discards reassembled frame.
					ret = NmParseFrame(myFrameParser, myRawFrame, i, 0, &myParsedFrame, NULL); 
					if(ret == ERROR_SUCCESS)
					{
						// Test to see if this frame passed our filter.
						BOOL passed = FALSE;
						ret = NmEvaluateFilter(myParsedFrame, myDisplayFilterID, &passed);
						if((ret == ERROR_SUCCESS) && (passed == TRUE))
						{							
							// know that strings are passed as word pointer to unicode string in the variant.
							GetFramePayload(myParsedFrame,myRawFrame,dumps);

							//WCHAR value[256];
							//ret = NmGetFieldValueString(myParsedFrame, myHTTPFieldID, 256, (LPWSTR)value);
							//if(ret == ERROR_SUCCESS)
							//{
							//	// Cast to WCHAR *
							//	wprintf(L"Frame %d: HTTP: %s\n", i+1, (LPWSTR)value);
							//}
						}
						// Release current parsed frame
						NmCloseHandle(myParsedFrame);
					}
					else
					{
						// Just print error but continue to loop
						wprintf(L"Error: 0x%X, trying to parse frame %d\n", ret, i+1);
					}

					// Release current raw frame
					NmCloseHandle(myRawFrame);
				}
				else
				{
					// Just print error but continue to loop
					wprintf(L"Errors getting raw frame %d\n", i+1);
				}
			}
		}

		NmCloseHandle(myFrameParser);
	}
	UnLoadNPL();
	return TRUE;
}